package com.ben.controller;

import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.ollama.OllamaChatModel;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

/**
 * @author huangyijun
 */
@RestController
@RequestMapping("/ai")
public class OllamaController {

  private final OllamaChatModel chatModel;

  public OllamaController(OllamaChatModel chatModel) {
    this.chatModel = chatModel;
  }

  // 普通同步回答
  @GetMapping("/chat")
  public String chat(@RequestParam String msg) {
    return chatModel.call(new Prompt(msg)).getResult().getOutput().getContent();
  }

  // Server-Sent Events 流式回答
  @GetMapping(value = "/stream", produces = "text/event-stream;charset=UTF-8")
  public Flux<String> stream(@RequestParam String msg) {
    return chatModel.stream(new Prompt(msg))
                    .map(resp -> resp.getResult().getOutput().getContent());
  }
}
